import java.util.*;

/**
 * Implementation of a priority queue for distinct elements.  Elements are
 * distinct if they are not equal, as determined by the  
 * {@link Object#equals(Object) equals(Object)} method.  The relative priority
 * of elements is determined by their natural ordering (see {@link Comparable}).
 * <br><br>
 * This priority queue is intended for use with elements whose natural ordering
 * is <i>not consistent with equals</i>.  The elements'
 * {@link Object#equals(Object) equals(Object)} and
 * {@link Comparable#compareTo(T) compareTo(T)} methods do not need to agree on
 * whether two elements are equal.  Two elements that are equal do not necessarily
 * have the same priority, and two elements with the same priority are not
 * necessarily equal.
 * <br><br>
 * If an attempt is made to add a non-distinct (a.k.a. duplicate) element to the
 * queue, the natural ordering of the two equal elements is used to determine
 * which one will remain in the queue and which one will be discarded.
 * The element with the lowest priority will remain in the queue.
 */
public class PriorityQueueSet<T extends Comparable<T>> {
	private ArrayList<T> heap;
	private Map<T,Integer> elementToHeapIndex;

	/**
	 * Initializes a new, empty <code>PriorityQueueSet</code> object.
	 */
	public PriorityQueueSet() {
		heap = new ArrayList<>();
		elementToHeapIndex = new HashMap<>();
	}

	/**
	 * Returns the number of elements in this queue.
	 *  
	 * @return The number of elements in this queue.
	 */
	public int size() {
		return heap.size();
	}

	/**
	 * Tests whether this queue contains no elements.
	 * 
	 * @return <code>true</code> if this queue contains no elements, or
	 * <code>false</code> otherwise.
	 */
	public boolean isEmpty() {
		return size() == 0;
	}

	/**
	 * Adds the specified element to this queue.  If the queue already contains the 
	 * specified element, the natural ordering of the two equal elements is used to
	 * determine which one remains in the queue and which one is discarded.  The
	 * element with the lowest priority remains in the queue.
	 * 
	 * @param element The element to add to this queue.
	 * @throws IllegalArgumentException If the specified element is <code>null</code>.
	 */
	public void add(T element) {
		if (element == null)
			throw new IllegalArgumentException();

		Integer index = elementToHeapIndex.get(element);
		if (index == null) {
			heap.add(element);
			bubbleUp(heap.size() - 1);
		} else if (element.compareTo(heap.get(index)) < 0) {
			heap.set(index, element);
			bubbleUp(index);
		}
	}

	/**
	 * Returns, but does not remove, the lowest priority element in this queue.
	 * 
	 * @return The lowest priority element in this queue.
	 * @throws IllegalStateException If the queue is empty.
	 */
	public T peek() {
		if (isEmpty())
			throw new IllegalStateException();
		
		return heap.get(0);
	}
	
	/**
	 * Removes and returns the lowest priority element in this queue.
	 * 
	 * @return The lowest priority element in this queue.
	 * @throws IllegalStateException If the queue is empty.
	 */
	public T remove() {
		if (heap.isEmpty())
			throw new IllegalStateException();

		T result = heap.get(0);
		elementToHeapIndex.remove(result);
		T last = heap.remove(heap.size() - 1);
		if (heap.size() > 0) {
			heap.set(0, last);
			bubbleDown();
		}
		return result;
	}
	
	private void bubbleUp(int index) {
		T tmp = heap.get(index);
		while (index > 0 && tmp.compareTo(heap.get((index - 1) / 2)) < 0) {
			heap.set(index, heap.get((index - 1) / 2));
			elementToHeapIndex.put(heap.get(index), index);
			index = (index - 1) / 2;
		};
		heap.set(index, tmp);
		elementToHeapIndex.put(tmp, index);
	}
	
	private void bubbleDown() {
		int index = 1;
		T tmp = heap.get(0);
		if (index + 1 < heap.size())
			if (heap.get(index + 1).compareTo(heap.get(index)) < 0)
				index++;
		while (index < heap.size() && heap.get(index).compareTo(tmp) < 0) {
			heap.set((index - 1) / 2, heap.get(index));
			elementToHeapIndex.put(heap.get((index - 1) / 2), (index - 1) / 2);
			index = 2 * index + 1;
			if (index + 1 < heap.size())
				if (heap.get(index + 1).compareTo(heap.get(index)) < 0)
					index++;
		};
		heap.set((index - 1) / 2, tmp);
		elementToHeapIndex.put(tmp, (index - 1) / 2);
	}
}
